package com.example.leetcode.tree;

import com.example.leetcode.tree.common.TreeNode;
import com.example.leetcode.tree.common.TreePrinter;
import com.example.leetcode.tree.common.TreeUtils;
import com.example.swing.BinaryTreeVisualizer;

import java.util.LinkedList;
import java.util.Queue;

public class leetcode623 {
    public static void main(String[] args) {
        Integer[] nums = {4, 2, 6, 3, 1, 5};
        int val = 1;
        int depth = 2;

        TreeNode root = TreeUtils.generateLeetcodeTree(nums);
        TreePrinter.printPrettyTree(root);
        BinaryTreeVisualizer.show(root);

        Solution solution = new Solution();
        TreeNode result = solution.addOneRow(root, val, depth);
        TreePrinter.printPrettyTree(result);
    }

    static class Solution {
        public TreeNode addOneRow(TreeNode root, int val, int depth) {
            if (depth == 1 || root == null) {
                TreeNode newRoot = new TreeNode(val);
                newRoot.left = root;
                return newRoot;
            }

            Queue<TreeNode> queue = new LinkedList<>();
            queue.add(root);
            int level = 0;

            while (!queue.isEmpty()) {
                int size = queue.size();
                level++;

                for (int i = 0; i < size; i++) {
                    TreeNode node = queue.poll();
                    assert node != null;
                    if (level == depth - 1) {
                        TreeNode newLeft = new TreeNode(val);
                        newLeft.left = node.left;
                        node.left = newLeft;

                        TreeNode newRight = new TreeNode(val);
                        newRight.right = node.right;
                        node.right = newRight;
                    } else {
                        if (node.left != null) {
                            queue.add(node.left);
                        }
                        if (node.right != null) {
                            queue.add(node.right);
                        }
                    }
                }

                if (level == depth - 1) {
                    break;
                }
            }

            return root;
        }
    }


}
